import socket
import struct
import binascii
import os,time

from cryptography import x509
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.backends import default_backend


def load_certificate():
    with open("../server.crt", "rb") as f:
        cert_der = f.read()
    return cert_der


def generate_ecdh_key():
    private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
    public_key = private_key.public_key()
    return private_key, public_key


def create_server_hello(client_version, client_random):
    # ServerHello参数
    version = client_version  # 使用客户端支持的最高版本
    random = struct.pack('>I', int(time.time())) + os.urandom(28)
    session_id = b'\x20' + os.urandom(31)  # 32字节会话ID
    cipher_suite = struct.pack('>H', 0xc030)  # 选择客户端支持的密码套件
    compression_method = b'\x00'

    # 扩展字段（示例包含ALPN）
    alpn_data = (
            b'\x00\x10' +  # 类型：ALPN
            struct.pack('>H', 5) +  # 长度
            b'\x02h2\x08http/1.1'
    )

    extensions = alpn_data
    extensions_len = struct.pack('>H', len(extensions))

    # 组合ServerHello消息
    server_hello = (
            b'\x02' +  # 消息类型：ServerHello
            struct.pack('>I', 32 + len(session_id) + 2 + 1 + 2 + len(extensions))[1:] +
            version +
            random +
            struct.pack('B', len(session_id)) +
            session_id +
            cipher_suite +
            compression_method +
            extensions_len +
            extensions
    )
    return server_hello


def create_certificate(cert_der):
    # 构造Certificate消息
    cert_length = struct.pack('>I', len(cert_der))[1:]
    return (
            b'\x0b' +  # 消息类型：Certificate
            cert_length +
            cert_der
    )


def create_server_key_exchange(private_key, public_key):
    # 构造ECDH ServerKeyExchange
    curve_params = struct.pack('>H', 0x001D)  # secp256r1

    # 序列化公钥
    public_bytes = public_key.public_bytes(
        encoding=serialization.Encoding.X962,
        format=serialization.PublicFormat.UncompressedPoint
    )[1:]  # 去除0x04前缀

    # 签名（示例使用临时私钥签名）
    signature = private_key.sign(
        b"server_key_exchange",
        ec.ECDSA(hashes.SHA256())
    )

    # 组合消息
    key_exchange = (
            b'\x0c' +  # 消息类型：ServerKeyExchange
            struct.pack('>I', 2 + len(public_bytes) + 4 + len(signature))[1:] +
            curve_params +
            struct.pack('>H', len(public_bytes)) +
            public_bytes +
            struct.pack('>I', len(signature))[1:] +
            signature
    )
    return key_exchange


def create_server_hello_done():
    return b'\x0e' + struct.pack('>I', 0)[1:]  # 消息类型：ServerHelloDone

# 封装握手消息到 TLS 记录层
def wrap_handshake(handshake_data):
    record_type = b'\x16'
    protocol_version = b'\x03\x03'
    record_length = struct.pack('>H', len(handshake_data))
    return record_type + protocol_version + record_length + handshake_data

# 初始化
cert_der = load_certificate()
private_key, public_key = generate_ecdh_key()

# 连接到客户端
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(('localhost', 8443))
s.listen(1)
conn, addr = s.accept()

# 接收ClientHello
client_hello = conn.recv(4096)

# 解析ClientHello获取版本信息
client_version = client_hello[4:6]

# 分四次发送服务器响应
server_hello = create_server_hello(client_version, client_hello[6:42])
conn.sendall(wrap_handshake(server_hello))
print("发送ServerHello：", binascii.hexlify(server_hello).decode())

certificate = create_certificate(cert_der)
server_key_exchange = create_server_key_exchange(private_key, public_key)
conn.sendall(wrap_handshake(certificate+server_key_exchange))
print("发送Certificate：", binascii.hexlify(certificate).decode())
print("发送ServerKeyExchange：", binascii.hexlify(server_key_exchange).decode())

server_hello_done = create_server_hello_done()
conn.sendall(wrap_handshake(server_hello_done))
print("发送ServerHelloDone：", binascii.hexlify(server_hello_done).decode())

conn.close()
s.close()